import os
from pdf2image import convert_from_path
import pytesseract
from PIL import Image
from typing import List, Tuple
from ..config import OCR_IMG_DIR, OCR_TEXT_DIR, OCR_LANG, OCR_CONFIG
from ..utils.logger import setup_logger

logger = setup_logger('ocr_service')

class OCRService:
    @staticmethod
    def convert_pdf_to_images(pdf_path: str) -> List[str]:
        """将PDF转换为图片"""
        try:
            # 创建输出目录
            os.makedirs(OCR_IMG_DIR, exist_ok=True)
            
            # 转换PDF为图片
            images = convert_from_path(pdf_path)
            image_paths = []
            
            # 保存图片
            for idx, image in enumerate(images):
                img_filename = f"{os.path.splitext(os.path.basename(pdf_path))[0]}_page_{idx}.png"
                img_path = os.path.join(OCR_IMG_DIR, img_filename)
                image.save(img_path, 'PNG')
                image_paths.append(img_path)
            
            logger.info(f"PDF转换完成，共{len(image_paths)}页")
            return image_paths
            
        except Exception as e:
            logger.error(f"PDF转换失败: {str(e)}")
            raise

    @staticmethod
    def perform_ocr(image_paths: List[str]) -> str:
        """对图片进行OCR识别"""
        try:
            # 创建输出目录
            os.makedirs(OCR_TEXT_DIR, exist_ok=True)
            
            # 对每页进行OCR
            all_text = []
            for img_path in image_paths:
                image = Image.open(img_path)
                text = pytesseract.image_to_string(
                    image,
                    lang=OCR_LANG,
                    config=OCR_CONFIG
                )
                all_text.append(text)
            
            # 合并所有文本
            combined_text = "\n\n".join(all_text)
            
            # 保存OCR结果
            text_filename = f"{os.path.splitext(os.path.basename(image_paths[0]))[0]}_ocr.txt"
            text_path = os.path.join(OCR_TEXT_DIR, text_filename)
            with open(text_path, 'w', encoding='utf-8') as f:
                f.write(combined_text)
            
            logger.info("OCR处理完成")
            return combined_text
            
        except Exception as e:
            logger.error(f"OCR处理失败: {str(e)}")
            raise

    @staticmethod
    def process_pdf(pdf_path: str) -> Tuple[str, List[str]]:
        """处理PDF文件：转换为图片并进行OCR"""
        try:
            # 转换为图片
            image_paths = OCRService.convert_pdf_to_images(pdf_path)
            
            # 进行OCR
            text = OCRService.perform_ocr(image_paths)
            
            return text, image_paths
            
        except Exception as e:
            logger.error(f"PDF处理失败: {str(e)}")
            raise 